import asyncio
import ssl
import websockets
import json
from aiortc_custom import RTCSessionDescription, RTCIceCandidate

class WebSocketSignaling:
    def __init__(self, server_url, client_id):
        self.server_url = server_url
        self.client_id = client_id
        self.websocket = None
        self.received_msgs = asyncio.Queue()
        self._receive_task = None
        self._closed = False
        
    async def connect(self):
        try:
            if self.server_url.startswith('wss://'):
                # SSL连接
                ssl_context = ssl.SSLContext()
                ssl_context.verify_mode = ssl.CERT_NONE
                self.websocket = await websockets.connect(
                    self.server_url, 
                    ssl=ssl_context,
                    ping_interval=20,  # 添加ping间隔
                    ping_timeout=10,   # ping超时
                    close_timeout=10   # 关闭超时
                )
            else:
                # 非SSL连接
                self.websocket = await websockets.connect(
                    self.server_url, 
                    ping_interval=20,  # 添加ping间隔
                    ping_timeout=10,   # ping超时
                    close_timeout=10   # 关闭超时
                )
            await self.websocket.send(self.client_id)
            self._receive_task = asyncio.create_task(self._receive_loop())
            print("WebSocket连接已建立")
        except Exception as e:
            print(f"WebSocket连接失败: {e}")
            raise
        
    async def _receive_loop(self):
        try:
            async for message in self.websocket:
                if self._closed:
                    break
                try:
                    data = json.loads(message)
                    await self.received_msgs.put(data)
                except json.JSONDecodeError as e:
                    print(f"JSON解析错误: {e}")
                except Exception as e:
                    print(f"处理消息时出错: {e}")
        except websockets.exceptions.ConnectionClosed as e:
            print(f"WebSocket连接已关闭: {e}")
        except Exception as e:
            print(f"WebSocket接收循环错误: {e}")
        finally:
            print("WebSocket接收循环已结束")
            
    async def close(self):
        """正确关闭WebSocket连接"""
        self._closed = True
        if self._receive_task:
            self._receive_task.cancel()
            try:
                await self._receive_task
            except asyncio.CancelledError:
                pass
        
        if self.websocket:
            try:
                await self.websocket.close(code=1000, reason="正常关闭")
                print("WebSocket连接已正确关闭")
            except Exception as e:
                print(f"关闭WebSocket时出错: {e}")
        
    async def send(self, desc):
        if self._closed or not self.websocket:
            raise Exception("WebSocket连接已关闭")
        message = {
            "type": "description",
            "description": {
                "type": desc.type,
                "sdp": desc.sdp
            },
            "target": self.peer_id
        }
        await self.websocket.send(json.dumps(message))
        
    async def send_candidate(self, candidate):
        if self._closed or not self.websocket:
            raise Exception("WebSocket连接已关闭")
        message = {
            "type": "candidate",
            "candidate": {
                "candidate": candidate.candidate,
                "sdpMLineIndex": candidate.sdpMLineIndex,
                "sdpMid": candidate.sdpMid
            },
            "target": self.peer_id
        }
        await self.websocket.send(json.dumps(message))

    async def receive(self):
        while not self._closed:
            try:
                data = await asyncio.wait_for(self.received_msgs.get(), timeout=30.0)
                if data["type"] == "description":
                    return RTCSessionDescription(
                        type=data["description"]["type"],
                        sdp=data["description"]["sdp"]
                    )
                elif data["type"] == "candidate":
                    candidate = data["candidate"]
                    return RTCIceCandidate(
                        candidate=candidate["candidate"],
                        sdpMLineIndex=candidate["sdpMLineIndex"],
                        sdpMid=candidate["sdpMid"]
                    )
            except asyncio.TimeoutError:
                if self._closed:
                    raise Exception("WebSocket连接已关闭")
                continue
            except Exception as e:
                print(f"接收消息时出错: {e}")
                if self._closed:
                    raise Exception("WebSocket连接已关闭")